
from mapper.DatabookMapper import DatabookMapper

databookMapper = DatabookMapper()

import os
import json
import hashlib
import urllib
columnKeyLocks = []
jsonDir = './json'

def _HttpRequest(url, headers={}, bodyData=None, method='GET'):
    UserAgent = 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_14_5) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/90.0.4430.212 Safari/537.36'
    AcceptLanguage = 'zh-CN,zh;q=0.9'
    Connection = 'keep-alive'

    if 'User-Agent' not in headers:
        headers['User-Agent'] = UserAgent

    if 'Accept-Language' not in headers:
        headers['Accept-Language'] = AcceptLanguage

    if 'Connection' not in headers:
        headers['Connection'] = Connection

    req = urllib.request.Request(url, headers=headers, data=bodyData, method=method)

    # 读取远程图片并保存到本地
    response = urllib.request.urlopen(req)
    return response

def getSuffix(data):
    if data[:4] == b'\x89PNG':
        # print(data[:20], filename)
        return 'png'
    elif data[:4] == b'RIFF' and data[8:15] == b'WEBPVP8':
        # print(data[:20], filename)
        return 'webp'
    elif data[:10] == b'\xff\xd8\xff\xe0\x00\x10JFIF' or (data[:4] == b'\xff\xd8\xff\xe1' and data[6:10] == b'Exif') or (data[:3] == b'\xff\xd8\xff' and ('jpg' in originalSuffix.lower() or 'jpeg' in originalSuffix.lower())):
        # print(data[:20], filename)
        return 'jpg'
    elif data[:6] == b'GIF89a' or data[:6] == b'GIF87a':
        # print(data[:20], filename)
        return 'gif'
    elif data[4:12] == b'ftypavif': # elif data[4:12] == b'ftypavif' and data[16:20] == b'avif': # b'\x00\x00\x00\x18ftypavif\x00\x00\x00\x00mif1'
        # print(data[:20], filename)
        return 'avif'
    elif (data[:4] == b'<svg' or data[:5] == b'<?xml'):
        # print(data[:20], filename)
        return 'svg'
    elif data[:2] == b'BM':
        # print(data[:20], filename)
        return 'bmp'
    else:
        return 'jpg'

class DatabookService:
    def __init__(self):
        print("DatabookService")

    def list (self, wname, bName):
        return databookMapper.list(wname, bName)

    def add (self, wname, bName, name):
        return databookMapper.add(wname, bName, name)

    def info (self, wname, bName, name):
        return databookMapper.info(wname, bName, name)

    def delete (self, wname, bName, name):
        return databookMapper.delete(wname, bName, name)

    def save (self, wname, bName, name, data):
        return databookMapper.save(wname, bName, name, data)

    def collectImage (self, wname, bName, name, columnKey):
        lockKey = "collectImage_{}_{}_{}_{}".format(wname, bName, name, columnKey)
        if lockKey in columnKeyLocks:
            return "任务进行中..."
        else:
            columnKeyLocks.append(lockKey)

            pathname = os.path.join(jsonDir, wname, bName, name)
            with open(pathname, mode="r", encoding="utf-8") as f:
                data = f.read()
                jsonData = json.loads(data)
                data = jsonData['data']

                for item in data:
                    if columnKey in item:
                        img_urls = item[columnKey]
                        index = 0
                        if type(img_urls) == list:
                            for img_url in img_urls:
                                if len(img_url) > 0:
                                    url_md5 = hashlib.md5(img_url.encode(encoding='UTF-8')).hexdigest()

                                    name_temp = img_url.replace('http://', '')
                                    name_temp = name_temp.replace('https://', '')
                                    name_temp_md5 = hashlib.md5(name_temp.encode(encoding='UTF-8')).hexdigest()

                                    saveFolder = os.path.join(jsonDir, wname, bName, name + '_save-image', columnKey)
                                    if os.path.exists(saveFolder) == False:
                                        os.makedirs(saveFolder)

                                    headers = {}
                                    if '360doc.com' in img_url or 'tiebapic.baidu.com' in img_url:
                                        headers['Referer'] = 'https://www.baidu.com'

                                    print(img_url)
                                    response = _HttpRequest(img_url, headers=headers)
                                    code = response.getcode() # 状态码

                                    if code == 200:
                                        data = response.read() # 数据

                                        suffix = getSuffix(data)
                                        img_filename = str(index) + '_' + url_md5 + '_' + name_temp_md5 + '.' + suffix
                                        img_path = os.path.join(saveFolder, img_filename)

                                        fd = open(img_path, mode="wb")
                                        fd.write(data)
                                        fd.close()

                                    index = index + 1
                        elif type(img_urls) == str:
                            img_url = img_urls
                            if len(img_url) > 0:
                                url_md5 = hashlib.md5(img_url.encode(encoding='UTF-8')).hexdigest()

                                name_temp = img_url.replace('http://', '')
                                name_temp = name_temp.replace('https://', '')
                                name_temp_md5 = hashlib.md5(name_temp.encode(encoding='UTF-8')).hexdigest()

                                saveFolder = os.path.join(jsonDir, wname, bName, name + '_save-image', columnKey)
                                if os.path.exists(saveFolder) == False:
                                    os.makedirs(saveFolder)

                                headers = {}
                                if '360doc.com' in img_url or 'tiebapic.baidu.com' in img_url:
                                    headers['Referer'] = 'https://www.baidu.com'

                                print(img_url)
                                response = _HttpRequest(img_url, headers=headers)
                                code = response.getcode() # 状态码

                                if code == 200:
                                    data = response.read() # 数据

                                    suffix = getSuffix(data)
                                    img_filename = str(index) + '_' + url_md5 + '_' + name_temp_md5 + '.' + suffix
                                    img_path = os.path.join(saveFolder, img_filename)

                                    fd = open(img_path, mode="wb")
                                    fd.write(data)
                                    fd.close()

            print('采集完成')
            columnKeyLocks.remove(lockKey)
            return columnKey